package stc.skymobi.netty.handler.codec.memcached;

import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import stc.skymobi.netty.initializer.MemcachedClientInitializer;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.ReentrantLock;
import java.util.zip.CRC32;

public class ChannelPool {
    private static Logger logger = LoggerFactory.getLogger(ChannelPool.class);
    // store instances of pools
    private static final ConcurrentMap<String, ChannelPool> pools = new ConcurrentHashMap<String, ChannelPool>();
    private static final NioEventLoopGroup NIO_EVENT_LOOP_GROUP = new NioEventLoopGroup();

    public static final int NATIVE_HASH = 0; // native String.hashCode();
    public static final int OLD_COMPAT_HASH = 1; // original compatibility
    // hashing algorithm (works with other clients)
    public static final int NEW_COMPAT_HASH = 2; // new CRC32 based
    // compatibility hashing algorithm (works with other clients)
    public static final int CONSISTENT_HASH = 3; // MD5 Based -- Stops

    private int hashingAlg = NATIVE_HASH; // default to using the native hash

    boolean initialized = false;
    private boolean aliveCheck = true; // default to check each connection
    // for being alive
    private boolean failover = true; // default to failover in event of cache
    private boolean failback = true; // only used if failover is also set ...
    // locks
    private final ReentrantLock initDeadLock = new ReentrantLock();

    private MemcachedClientInitializer memcachedClientPipelineFactory;

    // avoid recurring construction
    private static ThreadLocal<MessageDigest> MD5 = new ThreadLocal<MessageDigest>() {
        @Override
        protected final MessageDigest initialValue() {
            try {
                return MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                logger.error("++++ no md5 algorithm found");
                throw new IllegalStateException("++++ no md5 algorythm found");
            }
        }
    };

    // list of all servers
    private String[] servers;
    private Integer[] weights;
    private Integer totalWeight = 0;

    private List<String> buckets;
    private TreeMap<Long, String> consistentBuckets;

    // map to hold all available sockets
    private Map<String, Channel> channelPool;

    /**
     * Initializes the pool.
     */
    public void initialize() {
        initDeadLock.lock();
        try {

            // if servers is not set, or it empty, then
            // throw a runtime exception
            if (servers == null || servers.length <= 0) {
                logger.error("++++ trying to initialize with no servers");
                throw new IllegalStateException("++++ trying to initialize with no servers");
            }
            // pools
            channelPool = new HashMap<String, Channel>(servers.length);
            // initalize our internal hashing structures
            if (this.hashingAlg == CONSISTENT_HASH)
                populateConsistentBuckets();
            else
                populateBuckets();

            // mark pool as initialized
            this.initialized = true;

        } finally {
            initDeadLock.unlock();
        }
    }

    /**
     * Returns appropriate Channel object given string cache key and optional
     * hashcode.
     * <p/>
     * Trys to get Channel from pool. Fails over to additional pools in event of
     * server failure.
     *
     * @param key      hashcode for cache key
     * @param hashCode if not null, then the int hashcode to use
     * @return Channel obj connected to server
     */
    public final Channel getChannel(String key, Integer hashCode) {

        if (!this.initialized) {
            logger.error("attempting to get Channel from uninitialized pool!");
            return null;
        }

        // if no servers return null
        int size = 0;
        if ((this.hashingAlg == CONSISTENT_HASH && consistentBuckets.size() == 0)
                || (buckets != null && (size = buckets.size()) == 0))
            return null;
        else if (size == 1) {
            Channel channel = (this.hashingAlg == CONSISTENT_HASH) ? getConnection(consistentBuckets
                    .get(consistentBuckets.firstKey())) : getConnection(buckets.get(0));

            return channel;
        }

        // from here on, we are working w/ multiple servers
        // keep trying different servers until we find one
        // making sure we only try each server one time
        Set<String> tryServers = new HashSet<String>(Arrays.asList(servers));
        // get initial bucket
        long bucket = getBucket(key, hashCode);
        String server = (this.hashingAlg == CONSISTENT_HASH) ? consistentBuckets.get(bucket) : buckets
                .get((int) bucket);
        while (!tryServers.isEmpty()) {
            // try to get socket from bucket
            Channel channel = getConnection(server);
            if (channel != null)
                return channel;

            // if we do not want to failover, then bail here
            if (!failover)
                return null;
            // log that we tried
            tryServers.remove(server);
            if (tryServers.isEmpty())
                break;
            // if we failed to get a socket from this server
            // then we try again by adding an incrementer to the
            // current key and then rehashing
            int rehashTries = 0;
            while (!tryServers.contains(server)) {
                String newKey = String.valueOf(rehashTries) + key;
                // String.format( "%s%s", rehashTries, key );
                bucket = getBucket(newKey, null);
                server = (this.hashingAlg == CONSISTENT_HASH) ? consistentBuckets.get(bucket) : buckets
                        .get((int) bucket);
                rehashTries++;
            }
        }
        return null;
    }

    /**
     * Returns a Channel object from the pool for the passed in host.
     * <p/>
     * Meant to be called from a more intelligent method<br/>
     * which handles choosing appropriate server<br/>
     * and failover.
     *
     * @param host host from which to retrieve object
     * @return Channel object or null if fail to retrieve one
     */
    public final Channel getConnection(String host) {
        if (!this.initialized) {
            logger.error("attempting to get Channel from uninitialized pool!");
            return null;
        }
        if (host == null)
            return null;
        // if we have items in the pool then we can return it
        Channel channel = channelPool.get(host);
        if (channel == null) {
            channel = createAndPoolChannel(host);
            if (channel == null)
                return null;
        } else if (aliveCheck && !channel.isActive()) {
            channel = createAndPoolChannel(host);
            if (channel == null)
                return null;
        }
        return channel;
    }


    /**
     * Creates a new Channel obj for the given server.
     * <p/>
     * If server fails to connect, then return null and do not try<br/>
     * again until a duration has passed. This duration will grow<br/>
     * by doubling after each failed attempt to connect.
     *
     * @param host host:port to connect to
     * @return Channel obj or null if failed to create
     */
    protected final synchronized Channel createAndPoolChannel(String host) {
        Bootstrap bootstrap = new Bootstrap();
        String[] ip = host.split(":");
        bootstrap.group(NIO_EVENT_LOOP_GROUP)
                .channel(NioSocketChannel.class)
                .remoteAddress(ip[0], Integer.parseInt(ip[1]))
                .channel(NioSocketChannel.class)
                .handler(memcachedClientPipelineFactory);
        ChannelFuture future = bootstrap.connect();
        logger.info("connection to {} opened!", host);

        Channel channel;
        try {
            channel = future.await().channel();
            channelPool.put(host, channel);
        } catch (Exception ex) {
            logger.error("++++ failed to get Channel obj for: " + host);
            channel = null;
        }

        return channel;
    }

    private void populateBuckets() {
        // store buckets in tree map
        buckets = new ArrayList<String>();
        for (int i = 0; i < servers.length; i++) {
            if (this.weights != null && this.weights.length > i) {
                for (int k = 0; k < this.weights[i]; k++) {
                    buckets.add(servers[i]);
                }
            } else {
                buckets.add(servers[i]);
            }

            // Create the initial connections
            Channel channel = createAndPoolChannel(servers[i]);
            if (channel == null) {
                logger.error("++++ failed to create connection to: " + servers[i]);
                break;
            }
        }
    }

    private void populateConsistentBuckets() {
        // store buckets in tree map
        consistentBuckets = new TreeMap<Long, String>();

        MessageDigest md5 = MD5.get();
        if (this.totalWeight <= 0 && this.weights != null) {
            for (Integer weight : this.weights) this.totalWeight += (weight == null) ? 1 : weight;
        } else if (this.weights == null) {
            this.totalWeight = this.servers.length;
        }

        for (int i = 0; i < servers.length; i++) {
            int thisWeight = 1;
            if (this.weights != null && this.weights[i] != null)
                thisWeight = this.weights[i];

            double factor = Math.floor(((double) (40 * this.servers.length * thisWeight)) / (double) this.totalWeight);

            for (long j = 0; j < factor; j++) {
                byte[] d = md5.digest((servers[i] + "-" + j).getBytes());
                for (int h = 0; h < 4; h++) {
                    Long k = ((long) (d[3 + h * 4] & 0xFF) << 24) | ((long) (d[2 + h * 4] & 0xFF) << 16)
                            | ((long) (d[1 + h * 4] & 0xFF) << 8) | ((long) (d[(h * 4)] & 0xFF));

                    consistentBuckets.put(k, servers[i]);
                }
            }

            Channel channel = createAndPoolChannel(servers[i]);
            if (channel == null) {
                logger.error("++++ failed to create connection to: " + servers[i]);
                break;
            }
        }
    }

    private long getBucket(String key, Integer hashCode) {
        long hc = getHash(key, hashCode);

        if (this.hashingAlg == CONSISTENT_HASH) {
            return findPointFor(hc);
        } else {
            long bucket = hc % buckets.size();
            if (bucket < 0)
                bucket *= -1;
            return bucket;
        }
    }

    /**
     * Gets the first available key equal or above the given one, if none found,
     * returns the first k in the bucket
     *
     * @param hv key
     * @return
     */
    private Long findPointFor(Long hv) {
        // this works in java 6, but still want to release support for java5
        // Long k = this.consistentBuckets.ceilingKey( hv );
        // return ( k == null ) ? this.consistentBuckets.firstKey() : k;

        SortedMap<Long, String> tmap = this.consistentBuckets.tailMap(hv);

        return (tmap.isEmpty()) ? this.consistentBuckets.firstKey() : tmap.firstKey();
    }

    /**
     * Returns a bucket to check for a given key.
     *
     * @param key String key cache is stored under
     * @return int bucket
     */
    private long getHash(String key, Integer hashCode) {

        if (hashCode != null) {
            if (hashingAlg == CONSISTENT_HASH)
                return hashCode.longValue() & 0xffffffffL;
            else
                return hashCode.longValue();
        } else {
            switch (hashingAlg) {
                case NATIVE_HASH:
                    return (long) key.hashCode();
                case OLD_COMPAT_HASH:
                    return origCompatHashingAlg(key);
                case NEW_COMPAT_HASH:
                    return newCompatHashingAlg(key);
                case CONSISTENT_HASH:
                    return md5HashingAlg(key);
                default:
                    // use the native hash as a default
                    hashingAlg = NATIVE_HASH;
                    return (long) key.hashCode();
            }
        }
    }

    /**
     * Internal private hashing method.
     * <p/>
     * This is the original hashing algorithm from other clients. Found to be
     * slow and have poor distribution.
     *
     * @param key String to hash
     * @return hashCode for this string using our own hashing algorithm
     */
    private static long origCompatHashingAlg(String key) {
        long hash = 0;
        char[] cArr = key.toCharArray();

        for (char aCArr : cArr) {
            hash = (hash * 33) + aCArr;
        }

        return hash;
    }

    /**
     * Internal private hashing method.
     * <p/>
     * This is the new hashing algorithm from other clients. Found to be fast
     * and have very good distribution.
     * <p/>
     * UPDATE: This is dog slow under java
     *
     * @param key
     * @return
     */
    private static long newCompatHashingAlg(String key) {
        CRC32 checksum = new CRC32();
        checksum.update(key.getBytes());
        long crc = checksum.getValue();
        return (crc >> 16) & 0x7fff;
    }

    /**
     * Internal private hashing method.
     * <p/>
     * MD5 based hash algorithm for use in the consistent hashing approach.
     *
     * @param key
     * @return
     */
    private static long md5HashingAlg(String key) {
        MessageDigest md5 = MD5.get();
        md5.reset();
        md5.update(key.getBytes());
        byte[] bKey = md5.digest();
        long res = ((long) (bKey[3] & 0xFF) << 24) | ((long) (bKey[2] & 0xFF) << 16) | ((long) (bKey[1] & 0xFF) << 8)
                | (long) (bKey[0] & 0xFF);
        return res;
    }

    /**
     * Factory to create/retrieve new pools given a unique poolName.
     *
     * @param poolName unique name of the pool
     * @return instance of SockIOPool
     */
    public static ChannelPool getInstance(String poolName) {
        synchronized (pools) {
            if (!pools.containsKey(poolName)) {
                throw new RuntimeException("please initialize a connection first");
            }
        }

        return pools.get(poolName);
    }

    public void setMemcachedClientPipelineFactory(MemcachedClientInitializer memcachedClientPipelineFactory) {
        this.memcachedClientPipelineFactory = memcachedClientPipelineFactory;
    }

    public String[] getServers() {
        return servers;
    }

    public void setServers(String[] servers) {
        this.servers = servers;
    }
}
